import os
import numpy
import torch
from torch.utils.data import Dataset

class CustomDataset(Dataset):
    '''
    实现自定义数据集加载
    '''
    def __init__(self, root, train=True):
        self.root = root
        self.train = train

    def __len__(self):
        if self.train:
            file_list = os.listdir(os.path.join(self.root, "Train_data"))
        else:
            file_list = os.listdir(os.path.join(self.root, "Test_data"))
        return len(file_list)

    def __getitem__(self, index):
        # 确定文件路径
        if self.train:
            data_file = os.path.join(self.root, "Train_data", "{}_data.npy".format(index))
            target_file = os.path.join(self.root, "Train_lab", "{}_lab.npy".format(index))
        else:
            data_file = os.path.join(self.root, "Test_data", "{}_data.npy".format(index))
            target_file = os.path.join(self.root, "Test_lab", "{}_lab.npy".format(index))
        
        data = torch.from_numpy(numpy.load(data_file))
        label = torch.from_numpy(numpy.load(target_file))

        return data, label
